import numpy as np
from scipy.spatial import ConvexHull

class Point:
    def __init__(self, x=0, y=0, z=0):
        self.x = x
        self.y = y
        self.z = z

    def __repr__(self):
        return f"Point({self.x}, {self.y}, {self.z})"

    def to_array(self):
        return np.array([self.x, self.y, self.z])

class Polygon:
    def __init__(self, vertices):
        self.vertices = vertices

    def get_coordinates(self):
        return np.array([v.to_array() for v in self.vertices])

def compute_volume(vertices):
    """计算多面体的体积"""
    if len(vertices) < 4:
        return 0.0
    hull = ConvexHull(vertices)
    return hull.volume

def intersection_volume(box1, box2):
    """计算两个3D边界框的交集体积"""
    coords1 = box1.get_coordinates()
    coords2 = box2.get_coordinates()
    
    # 计算重叠部分的最小和最大坐标
    min_corner = np.maximum(coords1.min(axis=0), coords2.min(axis=0))
    max_corner = np.minimum(coords1.max(axis=0), coords2.max(axis=0))

    # 检查是否有重叠
    if np.any(min_corner >= max_corner):
        return 0.0

    # 计算交集的体积
    overlap_dimensions = max_corner - min_corner
    return np.prod(overlap_dimensions)

def iou_3d(box1, box2):
    """计算两个3D边界框的IoU"""
    volume1 = compute_volume(box1.get_coordinates())
    volume2 = compute_volume(box2.get_coordinates())
    inter_volume = intersection_volume(box1, box2)

    union_volume = volume1 + volume2 - inter_volume
    return inter_volume / union_volume if union_volume > 0 else 0

if __name__ == '__main__':
    # 定义两个3D边界框
    box1_vertices = [
        Point(0, 0, 0), Point(1, 0, 0), Point(1, 1, 0), Point(0, 1, 0),
        Point(0, 0, 1), Point(1, 0, 1), Point(1, 1, 1), Point(0, 1, 1)
    ]
    box2_vertices = [
        Point(0.5, 0.5, 0.5), Point(1.5, 0.5, 0.5), Point(1.5, 1.5, 0.5), Point(0.5, 1.5, 0.5),
        Point(0.5, 0.5, 1.5), Point(1.5, 0.5, 1.5), Point(1.5, 1.5, 1.5), Point(0.5, 1.5, 1.5)
    ]

    box1 = Polygon(box1_vertices)
    box2 = Polygon(box2_vertices)

    # 计算IoU
    iou = iou_3d(box1, box2)
    print(f"3D IoU: {iou:.4f}")
